import warnings
import pandas as pd
import datetime
import numpy as np

warnings.filterwarnings('ignore')

TRAIN_BANK_PATH = "./dataset/train_dataset/train_public.csv"
TRAIN_INTERNET_PATH = "./dataset/train_dataset/train_internet.csv"


def dataWashing():
    train_bank = pd.read_csv(TRAIN_BANK_PATH)
    train_internet = pd.read_csv(TRAIN_INTERNET_PATH)
    common_cols = []
    for col in train_bank.columns:
        if col in train_internet.columns:
            common_cols.append(col)
        else:
            continue

    train_bank['f0'] = train_bank['f0'].fillna(train_bank['f0'].mean()).to_frame()
    train_bank['f1'] = train_bank['f1'].fillna(train_bank['f1'].mean()).to_frame()
    train_bank['f2'] = train_bank['f2'].fillna(train_bank['f2'].mean()).to_frame()
    train_bank['f3'] = train_bank['f3'].fillna(train_bank['f3'].mean()).to_frame()
    train_bank['f4'] = train_bank['f4'].fillna(train_bank['f4'].mean()).to_frame()

    train_internet['f0'] = train_internet['f0'].fillna(train_internet['f0'].mean()).to_frame()
    train_internet['f1'] = train_internet['f1'].fillna(train_internet['f1'].mean()).to_frame()
    train_internet['f2'] = train_internet['f2'].fillna(train_internet['f2'].mean()).to_frame()
    train_internet['f3'] = train_internet['f3'].fillna(train_internet['f3'].mean()).to_frame()
    train_internet['f4'] = train_internet['f4'].fillna(train_internet['f4'].mean()).to_frame()

    train_internet_data = train_internet[common_cols]
    train_bank_data = train_bank[common_cols]

    # 日期类型：issueDate，earliesCreditLine
    # 转换为pandas中的日期类型
    train_internet_data['issue_date'] = pd.to_datetime(train_internet_data['issue_date'])
    # 提取多尺度特征
    train_internet_data['issue_date_y'] = train_internet_data['issue_date'].dt.year
    train_internet_data['issue_date_m'] = train_internet_data['issue_date'].dt.month
    # 提取时间diff
    # 设置初始的时间
    base_time = datetime.datetime.strptime('2007-06-01', '%Y-%m-%d')
    # 转换为天为单位
    train_internet_data['issue_date_diff'] = train_internet_data['issue_date'].apply(lambda x: x - base_time).dt.days
    train_internet_data[['issue_date', 'issue_date_y', 'issue_date_m', 'issue_date_diff']]
    train_internet_data.drop('issue_date', axis=1, inplace=True)

    # 日期类型：issueDate，earliesCreditLine
    # 转换为pandas中的日期类型
    train_bank_data['issue_date'] = pd.to_datetime(train_bank_data['issue_date'])
    # 提取多尺度特征
    train_bank_data['issue_date_y'] = train_bank_data['issue_date'].dt.year
    train_bank_data['issue_date_m'] = train_bank_data['issue_date'].dt.month
    # 提取时间diff
    # 设置初始的时间
    base_time = datetime.datetime.strptime('2007-06-01', '%Y-%m-%d')
    # 转换为天为单位
    train_bank_data['issue_date_diff'] = train_bank_data['issue_date'].apply(lambda x: x - base_time).dt.days
    train_bank_data[['issue_date', 'issue_date_y', 'issue_date_m', 'issue_date_diff']]
    train_bank_data.drop('issue_date', axis=1, inplace=True)

    employer_type = train_internet_data['employer_type'].value_counts().index
    industry = train_internet_data['industry'].value_counts().index

    emp_type_dict = dict(zip(employer_type, [0, 1, 2, 3, 4, 5]))
    industry_dict = dict(zip(industry, [i for i in range(15)]))

    train_internet_data['work_year'].fillna('10+ years', inplace=True)
    train_bank_data['work_year'].fillna('10+ years', inplace=True)

    work_year_map = {'10+ years': 10, '2 years': 2, '< 1 year': 0, '3 years': 3, '1 year': 1,
                     '5 years': 5, '4 years': 4, '6 years': 6, '8 years': 8, '7 years': 7, '9 years': 9}
    train_internet_data['work_year'] = train_internet_data['work_year'].map(work_year_map)
    train_bank_data['work_year'] = train_bank_data['work_year'].map(work_year_map)

    train_internet_data['class'] = train_internet_data['class'].map(
        {'A': 0, 'B': 1, 'C': 2, 'D': 3, 'E': 4, 'F': 5, 'G': 6})
    train_bank_data['class'] = train_bank_data['class'].map({'A': 0, 'B': 1, 'C': 2, 'D': 3, 'E': 4, 'F': 5, 'G': 6})

    train_internet_data['employer_type'] = train_internet_data['employer_type'].map(emp_type_dict)
    train_bank_data['employer_type'] = train_bank_data['employer_type'].map(emp_type_dict)

    is_default = train_internet_data['is_default']
    train_internet_data.pop('is_default')
    train_internet_data.insert(37, 'is_default', is_default)
    is_default = train_bank_data['is_default']
    train_bank_data.pop('is_default')
    train_bank_data.insert(37, 'is_default', is_default)

    train_internet_data['industry'] = train_internet_data['industry'].map(industry_dict)
    train_bank_data['industry'] = train_bank_data['industry'].map(industry_dict)
    train_bank_data = train_bank_data.drop('earlies_credit_mon', axis=1)
    train_bank_data = train_bank_data.drop('class', axis=1)
    train_bank_data = train_bank_data.drop('employer_type', axis=1)
    train_bank_data = train_bank_data.drop('industry', axis=1)

    train_internet_data = train_internet_data.drop('earlies_credit_mon', axis=1)
    train_internet_data = train_internet_data.drop('class', axis=1)
    train_internet_data = train_internet_data.drop('employer_type', axis=1)
    train_internet_data = train_internet_data.drop('industry', axis=1)

    return train_internet_data, train_bank_data


if __name__ == '__main__':
    _, train_bank_data=dataWashing()
    person_data = np.array(train_bank_data)
    x = person_data[:, :-1].astype(np.float64)
    print(x)